import os, re, cgi, json, time, hashlib, shutil, mimetypes, urllib.parse
from http.server import BaseHTTPRequestHandler, HTTPServer

HOST = "127.0.0.1"
PORT = 8988

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
UPLOAD_DIR = os.path.join(BASE_DIR, "uploads")
DOWNLOAD_DIR = os.path.join(BASE_DIR, "downloads")
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

if not mimetypes.inited:
    mimetypes.init()
EXT_MAP = mimetypes.types_map.copy()
EXT_MAP.update({"": "application/octet-stream", ".py": "text/plain", ".c": "text/plain", ".h": "text/plain"})

def guess_type(path: str):
    _, ext = os.path.splitext(path)
    return EXT_MAP.get(ext, EXT_MAP.get(ext.lower(), EXT_MAP[""]))

def safe_join(base: str, url_path: str) -> str:
    """Map URL to filesystem path inside `base` (prevent path traversal)."""
    # 允许 /downloads/xxx 前缀，也允许直接 /xxx
    p = url_path.split("?", 1)[0].split("#", 1)[0]
    if p.startswith("/downloads/"):
        p = p[len("/downloads/"):]
    p = urllib.parse.unquote(p).lstrip("/")

    parts = []
    for seg in p.split("/"):
        if not seg or seg in (".", ".."):
            continue
        parts.append(seg)
    full = os.path.join(base, *parts) if parts else base
    return full

def parse_range(range_header: str, size: int):
    """Return (start,end) or None."""
    if not range_header:
        return None
    m = re.match(r"bytes=(\d*)-(\d*)", range_header.strip())
    if not m:
        return None
    s, e = m.group(1), m.group(2)
    if s == "" and e == "":
        return None
    if s == "":
        # suffix-length
        suf = int(e)
        if suf <= 0: return None
        start = max(0, size - suf)
        end = size - 1
    else:
        start = int(s)
        end = size - 1 if e == "" else int(e)
    if start < 0 or start >= size or end < start:
        return None
    return (start, min(end, size - 1))

class Handler(BaseHTTPRequestHandler):
    # prettier logs
    def log_message(self, fmt, *args):
        print("[%s] %s" % (self.log_date_time_string(), fmt % args))

    def _common_headers(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Accept-Ranges", "bytes")

    # ---------- HEAD (downloads/) ----------
    def do_HEAD(self):
        path = safe_join(DOWNLOAD_DIR, self.path)
        if os.path.isdir(path):  # HEAD / -> 404
            self.send_error(404, "File not found")
            return
        if not os.path.isfile(path):
            self.send_error(404, "File not found")
            return

        size = os.path.getsize(path)
        ctype = guess_type(path)
        self.send_response(200)
        self._common_headers()
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(size))
        self.send_header("Last-Modified", self.date_time_string(os.path.getmtime(path)))
        self.end_headers()
        self.log_message("HEAD 200 file=%s len=%d ip=%s ua=%s",
                         os.path.basename(path), size, self.client_address[0], self.headers.get("User-Agent","-"))

    # ---------- GET (downloads/) ----------
    def do_GET(self):
        path = safe_join(DOWNLOAD_DIR, self.path)
        if os.path.isdir(path):
            self.send_error(404, "File not found")
            return
        if not os.path.isfile(path):
            self.send_error(404, "File not found")
            return

        size = os.path.getsize(path)
        ctype = guess_type(path)
        rng = parse_range(self.headers.get("Range"), size)
        t0 = time.time()
        try:
            with open(path, "rb") as f:
                if rng:
                    start, end = rng
                    length = end - start + 1
                    f.seek(start)
                    self.send_response(206)
                    self._common_headers()
                    self.send_header("Content-Type", ctype)
                    self.send_header("Content-Length", str(length))
                    self.send_header("Content-Range", f"bytes {start}-{end}/{size}")
                    self.send_header("Last-Modified", self.date_time_string(os.path.getmtime(path)))
                    self.end_headers()
                    # stream
                    remain = length
                    buf = 64 * 1024
                    while remain > 0:
                        chunk = f.read(min(buf, remain))
                        if not chunk: break
                        self.wfile.write(chunk)
                        remain -= len(chunk)
                    self.log_message("GET 206 file=%s range=%d-%d len=%d time=%.3fs ip=%s ua=%s",
                                     os.path.basename(path), start, end, length, time.time()-t0,
                                     self.client_address[0], self.headers.get("User-Agent","-"))
                else:
                    self.send_response(200)
                    self._common_headers()
                    self.send_header("Content-Type", ctype)
                    self.send_header("Content-Length", str(size))
                    self.send_header("Last-Modified", self.date_time_string(os.path.getmtime(path)))
                    self.end_headers()
                    shutil.copyfileobj(f, self.wfile)
                    self.log_message("GET 200 file=%s len=%d time=%.3fs ip=%s ua=%s",
                                     os.path.basename(path), size, time.time()-t0,
                                     self.client_address[0], self.headers.get("User-Agent","-"))
        except BrokenPipeError:
            self.log_message("GET aborted by client: %s", os.path.basename(path))

    def do_OPTIONS(self):
        self.send_response(204)
        self._common_headers()
        self.send_header("Access-Control-Allow-Methods", "GET,POST,HEAD,OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type,Range")
        self.end_headers()

    # ---------- POST (upload to uploads/) ----------
    def do_POST(self):
        t0 = time.time()
        if self.headers.get("Transfer-Encoding","").lower() == "chunked":
            self.send_response(411)
            self._common_headers()
            self.send_header("Content-Type","application/json; charset=utf-8")
            self.end_headers()
            self.wfile.write(json.dumps({"ok":False,"err":"chunked not supported"}).encode("utf-8"))
            return
        try:
            form = cgi.FieldStorage(
                fp=self.rfile, headers=self.headers,
                environ={
                    "REQUEST_METHOD":"POST",
                    "CONTENT_TYPE": self.headers.get("Content-Type"),
                    "CONTENT_LENGTH": self.headers.get("Content-Length","0"),
                }
            )
            if "file_data" not in form:
                raise ValueError("missing field: file_data")

            desc = form["desc"].value if "desc" in form else ""
            raw = form["file_data"].filename or "unnamed.bin"
            name = os.path.basename(raw)
            data = form["file_data"].value
            size = len(data)
            md5  = hashlib.md5(data).hexdigest()

            save_path = os.path.join(UPLOAD_DIR, name)
            with open(save_path, "wb") as f:
                f.write(data)

            host_hdr = self.headers.get("Host", f"{HOST}:{PORT}")
            url = f"http://{host_hdr}/downloads/{urllib.parse.quote(name)}"  # 可直接下载该文件（若手动放到 downloads 里）
            dt = time.time() - t0

            self.log_message("UPLOAD ok ip=%s ua=%s file=%s size=%d md5=%s saved=%s time=%.3fs",
                             self.client_address[0], self.headers.get("User-Agent","-"),
                             name, size, md5, save_path, dt)

            payload = {
                "ok": True,
                "file": name,
                "size": size,
                "desc": desc,
                "md5": md5,
                "saved_path": save_path,
                "url_hint": url,  # 提示：若你把同名文件放到 downloads/ 可用此 URL 下载
                "server_time": self.date_time_string(),
                "elapsed_sec": round(dt, 3),
            }
            self.send_response(201)
            self._common_headers()
            self.send_header("Content-Type","application/json; charset=utf-8")
            self.end_headers()
            self.wfile.write(json.dumps(payload, ensure_ascii=False).encode("utf-8"))
        except Exception as e:
            self.log_message("UPLOAD fail: %s", e)
            self.send_response(400)
            self._common_headers()
            self.send_header("Content-Type","application/json; charset=utf-8")
            self.end_headers()
            self.wfile.write(json.dumps({"ok":False,"err":str(e)}).encode("utf-8"))

def main():
    print(f"Server  : http://{HOST}:{PORT}")
    print(f"BASE    : {BASE_DIR}")
    print(f"UPLOADS : {UPLOAD_DIR}")
    print(f"DOWNLOAD: {DOWNLOAD_DIR}")
    httpd = HTTPServer((HOST, PORT), Handler)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down...")
    finally:
        httpd.server_close()

if __name__ == "__main__":
    main()
